import numpy as np
import sys
import time
from mpi4py import MPI
from mpi4py.MPI import ANY_SOURCE


# define the f(x)
def f(x):
    return x ** 2


# define integral function
def integrateRange(a, b, n):
    """
    Numerically integrate with the trapezoid rule on the interval from
    a to b with n trapezoids.
    """
    sum_f = -(f(a) + f(b)) / 2
    # sum of all linespace
    for x in np.linspace(a, b, n):
        sum_f = sum_f + f(x)
    integral = sum_f * (b - a) / n
    return integral


# get process rank and size
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

# takes in command-line arguments [a, b, n]
# a = float(sys.argv[1])
# b = float(sys.argv[2])
# n = int(sys.argv[3])
a = 0
b = 1
n = 100000

# local start time
start_time = time.perf_counter()

# init
integral = np.zeros(1)
recv_buffer = np.zeros(1)

# define local variables
h = (b - a) / n

# load balancing
local_n = int(n / size)
remainder = local_n % size
if rank < remainder:
    local_n += 1
local_a = a + rank * local_n * h
local_b = local_a + local_n * h

# integrate
integral[0] = integrateRange(local_a, local_b, local_n)
# integral[0] = rank

# Get levels from different processes
if rank == 0:
    total = integral[0]
    for i in range(1, size):
        comm.Recv(recv_buffer, ANY_SOURCE)
        total += recv_buffer[0]
else:
    comm.Send(integral, dest=0)

# local stop time
stop_time = time.perf_counter()
if comm.rank == 0:
    # print(ranks, sum(ranks))
    print(
        f"With n = {n} trapezoids, our estimate of the integral from {a} to {b} is {total}")
    print(f"Time cost: {stop_time-start_time:.8f}")
